import json

import networkx as nx

from model import Discipline


def generate_coordinates(data: list[Discipline], force_iter=1000,
                         canvas_size=1000, ) -> dict:
    """生成网络布局坐标"""
    g = nx.Graph()
    for node in data:
        g.add_node(node.name)

    for node in data:
        if node.super_disciplines:
            g.add_edge(node.super_disciplines[-1], node.name)

    cord = nx.spring_layout(g, iterations=force_iter)

    cord = {k: [v[0] * canvas_size, v[1] * canvas_size] for k, v in cord.items()}

    return cord


if __name__ == '__main__':
    with open("../assets/base_data.json", 'r') as f:
        base_data = [Discipline.parse_raw(d) for d in json.load(f)]

    # 生成坐标，设置画布大小为1000x1000，节点间距为30-100
    coordinates = generate_coordinates(
        base_data,
        force_iter=500,
        canvas_size=1000,
    )

    with open("../assets/coordinate_data.json", "w", encoding="utf-8") as f:
        json.dump(coordinates, f, ensure_ascii=False, indent=2)
